#!/usr/bin/env python

###############################################################################
#
# MODULE:       m.fuzzy.sensitivity
# AUTHOR:       Thomas Leppelt
# PURPOSE:      Compute sensitivity anaylsis for fuzzy models in ranges of given 
#               raster cell maps or by direct input and plotting the output in R.
#
###############################################################################

#%module
#% description: Compute sensitivity anaylsis for fuzzy models.
#% keywords: Miscellaneous
#% keywords: Fuzzy model,TAG2E, sensitivity analysis 
#%end
#%option G_OPT_F_INPUT
#% description: List of input raster maps representing driving parameters of a fuzzy model.
#% required : yes
#% multiple: yes
#% guisection: Files & format
#%end
#%option G_OPT_F_OUTPUT
#% description: Basename of output vector map containing model results of sensitivity analysis.
#% required : yes
#% guisection: Files & format
#%end
#%option
#% key: parameter
#% type : string
#% description: List of XML fuzzy inference parameter file names.
#% required : no
#% multiple: yes
#% guisection: Files & format
#%end
#%option
#% key: basename
#% type : string
#% description: Basename of vector files for calculated fuzzy models.
#% required : yes
#% multiple: no
#% guisection: Files & format
#%end
#%option
#% key: file
#% type : string
#% description: Text file containing XML fuzzy inference parameter file names as list.
#% required : no
#% multiple: no
#% guisection: Files & format
#%end
#%option
#% key: driver
#% type : string
#% description: Name of the model drivers in the order of the raster map list that are used also in the XML fuzzy inference scheme file.
#% required : yes
#% multiple: yes
#% guisection: Files & format
#%end
#%option
#% key: breaks
#% type : integer
#% description: Number of equally sized breaks per parameter.
#% required : yes
#% multiple: no
#% answer: 10
#% guisection: Files & format
#%end
#%flag
#% key: p
#% description: Plot the sensitivity analysis result as 2D/3D vector graphic in R (Packages: scatterplot3d & SDMTools are required).
#%end

import grass.script as grass
import grass.pygrass.modules as pygrass
from grass.pygrass.vector import VectorTopo
from grass.pygrass.vector.geometry import Point
from grass.pygrass.vector.table import Link
import itertools as it
from subprocess import PIPE
import numpy as np
import os
import sys
import csv
import subprocess
import tempfile

def Get_Parameter_Range(parameters, drivers, breaks = 10):
    """!This function calculates parameter ranges from given input raster cell maps 
        for given driving parameter list.

        @param parameter   List of raster cell maps representing driving 
                           parameters of a fuzzy model.
        @param drivers     List of driver names in same order as given corresponding 
                           raster maps.
        @param breaks      Number of steps created in between ranges of parameters.
        
        @return Dictionary with parameter ranges.
    """
    # Check if parameter and drivers have same length.
    if len(parameters) != len(drivers):
        grass.fatal("Length of parameters (%s) and fuzzy model drivers (%s) does not match." 
                        %(parameters, drivers))

    # Check driving parameter list. 
    if isinstance(parameters, list) == False:
        grass.fatal("Expected parameter list object. Got:  %s." %(parameters))
        
    # Build pyGrass module for r.info.
    mod_rinfo = pygrass.Module('r.info')
    mod_rinfo.flags["r"].value = True

    para_dict = {}
    range_dict = {}
    # Number of generated values.
    n = breaks

    # Loop over parameter list and build dictionary with range values of parameter.
    for i in range(len(parameters)):
        para = parameters[i]
        mod_rinfo(map = str(para), stdout_= PIPE)
        para_range = mod_rinfo.outputs.stdout.strip().split('\n')
        range_dict[drivers[i]] = [float(para_range[0][4:]),float(para_range[1][4:])]

    # Create parameter lists with n steps in given ranges. 
    for i in range(len(parameters)):
        para = drivers[i]
        para_min = range_dict[para][0]
        para_max = range_dict[para][1]
        para_range = para_max - para_min
        step = para_range / n
        para_vect = np.arange(para_min, para_max, step)
        para_list = list(para_vect) + [para_max]
        para_dict[para] = para_list

    # Return dictionary with range values for given parameters.
    return(para_dict)

###############################################################################
def Run_Fuzzy_Model(name, model_list, drivers, para_ranges, basename):
    """!This function generates input data and calculates fuzzy models for given 
        input XML fuzzy model files.

        @param name         Name of output vector file with data in given ranges.
        @param model_list   List of fuzzy model XML representations.
        @param drivers      List of driver names that are used in the fuzzy model 
                            XML representation.
        @param para_ranges  Dictionary with parameter ranges.
        @param basename     Basename string to name calculated fuzzy models.
                        .
        @return Dictionary with parameter ranges.
    """

    # Build pyGrass module for fuzzy model application.
    mod_vfuzzy = pygrass.Module('v.fuzzy.model')
    mod_vfuzzy.flags.overwrite = grass.overwrite()
    
    # Fuzzy model sensitivity analysis.

    # Create input vector files.
    model_data = VectorTopo(name)
    model_data.overwrite = grass.overwrite()
    
    # Generate attribute data.
    data_product = [x for x in apply(it.product, para_ranges.values())]
    
    # Set the layer name.
    model_data.layer = 1
    # Create a Link object.
    link = Link(layer=model_data.layer, name=model_data.name, table=model_data.name, 
                key='cat', database='$GISDBASE/$LOCATION_NAME/$MAPSET/sqlite/sqlite.db',
                driver='sqlite')
                
    # Define columns and create column list.
    cols = [(u'cat',        'INTEGER PRIMARY KEY')]
    for col in para_ranges.keys():
        cols.append((u'%s'%(col), 'DOUBLE PRECISION'))
    cols.append((u'result_mean', 'DOUBLE PRECISION'))
    
    # Apply fuzzy model on generated data.
    count = 0
    result_all = []
    for xml in model_list:
        out_name = basename + '_' + str(count)
        # Calculate fuzzy model values.
        mod_vfuzzy(input = name, parameter = xml, output = out_name)
        # Read result values from vector table.
        result = VectorTopo(out_name)
        #result.build()
        result.open('r')
        link = result.dblinks[0]
        conn = link.connection()
        cur = conn.cursor()
        cur.execute("SELECT result FROM %s" % result.name)
        print("SELECT result FROM %s" % result.name)
        result_vec = cur.fetchall()
        # Store results in table.
        if count == 0:
            result_all = result_vec
        else:
            result_all = [x+(y) for x,y in zip(result_all, result_vec)]
        cur.close()
        conn.close()
        result.close()
        count= count + 1

    result_mean = [np.mean(i) for i in result_all]
    result_table = [x+(y,) for x,y in zip(data_product, result_mean)]
    
    # Open vector object in write mode.
    model_data.open('w', tab_cols=cols)

    # Add link to vector map.
    model_data.dblinks.add(link)

    # Write vector poitns and parameter data to table.
    for i in range(len(data_product)):
        point_i = Point(i, i)
        p_attr = data_product[i]
        model_data.write(point_i, (str(p_attr[0]),str(p_attr[1]),str(p_attr[2]), 
                            str(result_mean[i]), ))
    
    # Save the changes and close the map.
    model_data.table.conn.commit()
    model_data.close()
    
    return(result_table)
 
def Plot_Fuzzy_Result(name, result, drivers):
    """!This function generates a csv file of the sensitivity analysis result 
        in temporary directory and creates a 2D/3D plot of the results.

    @param name    Name of the export csv file for input data.
    @param result  Name of list containing results.
    @param drivers List of drivers for the fuzzy model.
                    .
    @return File containing plotted results.
    """

    csv_dir = tempfile.gettempdir()
    csv_name = str(name + '.csv')
    # Export result table as csv.
    o = open(str(csv_dir + '/' + csv_name), "w+")
    w = csv.writer(o)
    cols = []
    for col in drivers:
        cols.append(col)
    cols.append('result_mean')
    w.writerow(cols)
    for row in result:
        w.writerow(row)
    print("Exporting " + str(csv_dir + '/' + csv_name))
    # Close csv file.
    o.close()
    
    # Execute R-script.
    RScript = """
    # Perform sensitivity analysis of fuzzy models for different land use categories.

    require(scatterplot3d)
    require(SDMTools)

    add_plots = FALSE

    # Export as pdf file.
    exportpdf <- paste("outputname", "_", Sys.Date(), '.pdf', sep = '')
    pdf(file = exportpdf)

    ###############################################################################
    # Fuzzy model sensitivity.

    # Model directory and file name.
    model_sensitive_dir = "temp_dir"
    model_sensitive_name = "fuzzy_result"

    # Import model sensitivity analysis data.
    modeldata <- read.csv(paste(model_sensitive_dir,'/',model_sensitive_name, sep=''))
    
    # Check number of columns.
    modellen <- length(modeldata)
    
    # Transform fluxes. 
    modeldata$result_mean <- exp(modeldata$result_mean) - 0.5

    # 3D plot.
    # Create a function to generate a continuous color palette.
    colPal <- colorRampPalette(c('yellow', 'green', 'cyan', 'blue', 'violet', 'red'))
    # This adds a column of color values based on the emission values.
    modeldata$col <- colPal(200)[as.numeric(cut(modeldata$result_mean,breaks = 200))]
    modeldata$cex <- as.numeric(cut(modeldata$result_mean,breaks = 10))/2
    
    xlim <- range(modeldata[,1])
    ylim <- range(modeldata[,2])
    zlim <- range(modeldata[,3])
    scatterplot3d(x = modeldata[,1], y = modeldata[,2], z = modeldata[,3], color = modeldata$col, 
        cex.symbols = modeldata$cex, pch = 16, xlim=xlim, ylim=ylim, zlim=zlim)
    pnts = (cbind(x =c(max(modeldata[,1]),max(modeldata[,1]) + 0.5,max(modeldata[,1]) + 
                    0.5,max(modeldata[,1])), y =c(max(modeldata[,3]) + 
                    7.5,max(modeldata[,3]) + 5.5,max(modeldata[,3]) + 
                    5.5,max(modeldata[,3]) + 2)))
    legend.gradient(pnts,cols=colPal(30),limits=round(range(modeldata$result_mean),2), 
        title='Flux')

    if (add_plots) {
        # Extract high fluxes.
        modeldata_high <- subset(modeldata, result_mean > 0.8)
        with(modeldata_high,scatterplot3d(ph,rr_annual,wt, color = col, cex.symbols = cex, 
        pch = 16, xlim=xlim, ylim=ylim, zlim=zlim))

        modeldata_low <- subset(modeldata, result_mean <= 0.8)
        with(modeldata_low,scatterplot3d(ph,rr_annual,wt, color = col, cex.symbols = cex, 
        pch = 16, xlim=xlim, ylim=ylim, zlim=zlim))
        }
    print("Plotting Result")
    ###############################################################################
    """
    RScript = RScript.replace("outputname", name)
    RScript = RScript.replace("temp_dir", csv_dir)
    RScript = RScript.replace("fuzzy_result", csv_name)
    inputlist = ["R", "--vanilla"]
    proc = subprocess.Popen(args=inputlist, stdin=subprocess.PIPE)
    proc.stdin.write(RScript)
    proc.communicate()

def main():
  
    input = options['input']
    output = options['output']
    parameter = options['parameter']
    basename = options['basename']
    filename = options['file']
    driver = options['driver']
    breaks = options['breaks']
    plot = flags['p']

    # Check if parameters and file name are both input.
    if parameter and filename:
        grass.fatal("Expected either parameter list or file with list of parameters. Got both.")

    # Convert inputs into correct format.
    input_list = input.split(',')
    driver_list = driver.split(',')
    breaks = float(breaks)
    if parameter:
        parameter_list = parameter.split(',')
    elif filename:
        parameter_list = []
        with open(filename, 'rb') as xmlfile:
            xmlreader = csv.reader(xmlfile, delimiter=',')
            for row in xmlreader:
                parameter_list = parameter_list + (row)

    # Calculate parameter ranges.
    para_dict = Get_Parameter_Range(input_list, driver_list, breaks)
    # Calculate fuzzy models and create output vector file. 
    result_list = Run_Fuzzy_Model(output, parameter_list, driver_list, 
                                    para_dict, basename)
    # Plotting the results in R.
    if plot:
        Plot_Fuzzy_Result(output, result_list, driver_list)
        
if __name__ == "__main__":
    if len(sys.argv) == 2 and sys.argv[1] == 'doctest':
        import doctest
        doctest.testmod()
        exit()
    options, flags = grass.parser()
    main()
